#!/bin/bash
# 配置参数 - 根据实际情况修改
APP_NAME="resource-occupier"
SERVICE_NAME="${APP_NAME}.service"
MONITOR_SCRIPT_PATH="$(pwd)/monitor.sh"  # 监控脚本的绝对路径
WORKING_DIR="$(pwd)"                     # 应用程序工作目录
USER=$(logname)                          # 运行程序的用户

# 检查是否以root权限运行
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "错误: 请使用root权限运行此脚本 (sudo ./resource_manager.sh [install|uninstall|start|stop|restart|status])"
        exit 1
    fi
}

# 检查监控脚本是否存在
check_monitor_script() {
    if [ ! -f "$MONITOR_SCRIPT_PATH" ]; then
        echo "错误: 监控脚本 $MONITOR_SCRIPT_PATH 不存在"
        exit 1
    fi
}

# 安装服务（添加到系统启动项）
install_service() {
    check_root
    check_monitor_script
    
    if [ -d "/run/systemd/system" ]; then
        echo "检测到systemd系统，安装服务..."
        
        # 创建systemd服务文件
        SERVICE_FILE="/etc/systemd/system/$SERVICE_NAME"
        cat > "$SERVICE_FILE" << EOF
[Unit]
Description=Resource Occupier Monitor Service
After=network.target

[Service]
User=$USER
WorkingDirectory=$WORKING_DIR
ExecStart=$MONITOR_SCRIPT_PATH
Restart=always
RestartSec=5
StandardOutput=append:$WORKING_DIR/startup.log
StandardError=append:$WORKING_DIR/startup_error.log

[Install]
WantedBy=multi-user.target
EOF

        # 启用并启动服务
        systemctl daemon-reload
        systemctl enable "$SERVICE_NAME"
        systemctl start "$SERVICE_NAME"
        
        # 检查服务状态
        echo "服务安装完成，当前状态："
        systemctl status "$SERVICE_NAME" --no-pager
    elif [ -f "/etc/init.d/functions" ]; then
        echo "检测到sysvinit系统，安装服务..."
        
        # 创建sysvinit脚本
        INIT_SCRIPT="/etc/init.d/$APP_NAME"
        cat > "$INIT_SCRIPT" << EOF
#!/bin/bash
# chkconfig: 345 99 10
# description: Resource Occupier Monitor

WORKING_DIR=$WORKING_DIR
MONITOR_SCRIPT=$MONITOR_SCRIPT_PATH
USER=$USER

start() {
    echo "启动$APP_NAME服务..."
    su - \$USER -c "cd \$WORKING_DIR && nohup \$MONITOR_SCRIPT &"
}

stop() {
    echo "停止$APP_NAME服务..."
    pkill -f "\$MONITOR_SCRIPT"
    pkill -f "ResourceOccupier"
}

restart() {
    stop
    sleep 2
    start
}

status() {
    if pgrep -f "\$MONITOR_SCRIPT" > /dev/null; then
        echo "$APP_NAME服务正在运行"
        return 0
    else
        echo "$APP_NAME服务未运行"
        return 1
    fi
}

case "\$1" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    restart)
        restart
        ;;
    status)
        status
        ;;
    *)
        echo "使用方法: \$0 {start|stop|restart|status}"
        exit 1
        ;;
esac
EOF

        # 设置权限并添加到启动项
        chmod +x "$INIT_SCRIPT"
        chkconfig --add "$APP_NAME"
        chkconfig "$APP_NAME" on
        service "$APP_NAME" start
        
        # 检查服务状态
        echo "服务安装完成，当前状态："
        service "$APP_NAME" status
    else
        echo "无法识别的系统初始化类型，无法安装服务"
        exit 1
    fi

    echo "安装完成！系统启动时将自动运行资源占用程序"
}

# 卸载服务（从系统启动项中移除）
uninstall_service() {
    check_root
    
    if [ -d "/run/systemd/system" ]; then
        echo "检测到systemd系统，卸载服务..."
        
        # 停止并禁用服务
        systemctl stop "$SERVICE_NAME" 2>/dev/null
        systemctl disable "$SERVICE_NAME" 2>/dev/null
        
        # 删除服务文件
        SERVICE_FILE="/etc/systemd/system/$SERVICE_NAME"
        if [ -f "$SERVICE_FILE" ]; then
            rm -f "$SERVICE_FILE"
        fi
        
        # 重载配置
        systemctl daemon-reload
        
        echo "服务已卸载"
    elif [ -f "/etc/init.d/functions" ]; then
        echo "检测到sysvinit系统，卸载服务..."
        
        # 停止服务
        service "$APP_NAME" stop 2>/dev/null
        
        # 移除开机启动
        chkconfig --del "$APP_NAME" 2>/dev/null
        
        # 删除服务脚本
        INIT_SCRIPT="/etc/init.d/$APP_NAME"
        if [ -f "$INIT_SCRIPT" ]; then
            rm -f "$INIT_SCRIPT"
        fi
        
        echo "服务已卸载"
    else
        echo "无法识别的系统初始化类型，无法卸载服务"
        exit 1
    fi
    
    # 清理残留进程
    echo "清理残留进程..."
    pkill -f "monitor.sh" 2>/dev/null
    pkill -f "ResourceOccupier" 2>/dev/null
    
    echo "卸载完成"
}

# 启动服务
start_service() {
    check_root
    
    if [ -d "/run/systemd/system" ]; then
        systemctl start "$SERVICE_NAME"
        systemctl status "$SERVICE_NAME" --no-pager
    else
        service "$APP_NAME" start
    fi
}

# 停止服务
stop_service() {
    check_root
    
    if [ -d "/run/systemd/system" ]; then
        systemctl stop "$SERVICE_NAME"
        systemctl status "$SERVICE_NAME" --no-pager
    else
        service "$APP_NAME" stop
    fi
    
    # 清理残留进程
    pkill -f "monitor.sh" 2>/dev/null
    pkill -f "ResourceOccupier" 2>/dev/null
}

# 重启服务
restart_service() {
    check_root
    
    if [ -d "/run/systemd/system" ]; then
        systemctl restart "$SERVICE_NAME"
        systemctl status "$SERVICE_NAME" --no-pager
    else
        service "$APP_NAME" restart
    fi
}

# 查看服务状态
status_service() {
    check_root
    
    if [ -d "/run/systemd/system" ]; then
        systemctl status "$SERVICE_NAME" --no-pager
    else
        service "$APP_NAME" status
    fi
}

# 显示帮助信息
show_help() {
    echo "资源占用程序管理脚本"
    echo "用法: sudo ./resource_manager.sh [命令]"
    echo "命令列表:"
    echo "  install    - 安装服务并添加到系统启动项"
    echo "  uninstall  - 卸载服务并从系统启动项中移除"
    echo "  start      - 启动服务"
    echo "  stop       - 停止服务"
    echo "  restart    - 重启服务"
    echo "  status     - 查看服务状态"
    echo "  help       - 显示帮助信息"
}

# 根据参数执行相应操作
case "$1" in
    install)
        install_service
        ;;
    uninstall)
        uninstall_service
        ;;
    start)
        start_service
        ;;
    stop)
        stop_service
        ;;
    restart)
        restart_service
        ;;
    status)
        status_service
        ;;
    help|*)
        show_help
        ;;
esac
